import os
from os.path import join
import numpy as np
from tqdm import tqdm
import open3d as o3d


def intial_position_distribution_visualization(dataset_dir, save_path, task_is_touch_point):
    positions = []
    for obj_name in tqdm(os.listdir(dataset_dir)):
        if not task_is_touch_point:
            obj_motion_dir = join(dataset_dir, obj_name, "h1_kinematic_motions")
            try:
                for motion_name in os.listdir(obj_motion_dir):
                    motion_fp = join(obj_motion_dir, motion_name)
                    position = np.load(motion_fp, allow_pickle=True)["arr_0"].item()["global_translations"][0]
                    positions.append(position)
            except:
                print("error:", obj_motion_dir)
        else:
            # for touchpoint
            for seq_name in os.listdir(join(dataset_dir, obj_name)):
                obj_motion_dir = join(dataset_dir, obj_name, seq_name, "h1_kinematic_motions")
                try:
                    for motion_name in os.listdir(obj_motion_dir):
                        motion_fp = join(obj_motion_dir, motion_name)
                        position = np.load(motion_fp, allow_pickle=True)["arr_0"].item()["global_translations"][0]
                        positions.append(position)
                except:
                    print("error:", obj_motion_dir)

    positions = np.float32(positions)
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(positions)
    o3d.io.write_point_cloud(save_path, pcd)


if __name__ == "__main__":
    intial_position_distribution_visualization("/media/liuyun/Seagate Basic/Humanoid_IL_Benchmark/retargeted_data/UniHSI_partnet_sit_chair", "./initial_positions_SC.ply", task_is_touch_point=False)
    intial_position_distribution_visualization("/media/liuyun/Seagate Basic/Humanoid_IL_Benchmark/retargeted_data/UniHSI_partnet_lie_bed_long", "./initial_positions_LB.ply", task_is_touch_point=False)
    intial_position_distribution_visualization("/media/liuyun/Seagate Basic/Humanoid_IL_Benchmark/retargeted_data/ROAM_sofa_sit_10000_optim_Batch1_examples", "./initial_positions_SS.ply", task_is_touch_point=False)
    intial_position_distribution_visualization("/media/liuyun/Seagate Basic/Humanoid_IL_Benchmark/retargeted_data/ROAM_sofa_lie_10000_optim_Batch1_examples", "./initial_positions_LS.ply", task_is_touch_point=False)
    intial_position_distribution_visualization("/media/liuyun/Seagate Basic/Humanoid_IL_Benchmark/retargeted_data/CORE4D_copyrot_touch_point", "./initial_positions_T.ply", task_is_touch_point=True)
    intial_position_distribution_visualization("/media/liuyun/Seagate Basic/Humanoid_IL_Benchmark/retargeted_data/CORE4D_lift_box_1018", "./initial_positions_L.ply", task_is_touch_point=False)
